import numpy as np
import re
from collections import defaultdict, Counter
import math
import random
from typing import List, Tuple, Dict, Set
import matplotlib.pyplot as plt

# 本地英文停用词列表
ENGLISH_STOP_WORDS = {
    'i', 'me', 'my', 'myself', 'we', 'our', 'ours', 'ourselves', 'you', "you're", 
    "you've", "you'll", "you'd", 'your', 'yours', 'yourself', 'yourselves', 'he', 
    'him', 'his', 'himself', 'she', "she's", 'her', 'hers', 'herself', 'it', 
    "it's", 'its', 'itself', 'they', 'them', 'their', 'theirs', 'themselves', 
    'what', 'which', 'who', 'whom', 'this', 'that', "that'll", 'these', 'those', 
    'am', 'is', 'are', 'was', 'were', 'be', 'been', 'being', 'have', 'has', 
    'had', 'having', 'do', 'does', 'did', 'doing', 'a', 'an', 'the', 'and', 
    'but', 'if', 'or', 'because', 'as', 'until', 'while', 'of', 'at', 'by', 
    'for', 'with', 'through', 'during', 'before', 'after', 'above', 'below', 
    'up', 'down', 'in', 'out', 'on', 'off', 'over', 'under', 'again', 'further', 
    'then', 'once', 'here', 'there', 'when', 'where', 'why', 'how', 'all', 
    'any', 'both', 'each', 'few', 'more', 'most', 'other', 'some', 'such', 
    'no', 'nor', 'not', 'only', 'own', 'same', 'so', 'than', 'too', 'very', 
    's', 't', 'can', 'will', 'just', 'don', "don't", 'should', "should've", 
    'now', 'd', 'll', 'm', 'o', 're', 've', 'y', 'ain', 'aren', "aren't", 
    'couldn', "couldn't", 'didn', "didn't", 'doesn', "doesn't", 'hadn', 
    "hadn't", 'hasn', "hasn't", 'haven', "haven't", 'isn', "isn't", 'ma', 
    'mightn', "mightn't", 'mustn', "mustn't", 'needn', "needn't", 'shan', 
    "shan't", 'shouldn', "shouldn't", 'wasn', "wasn't", 'weren', "weren't", 
    'won', "won't", 'wouldn', "wouldn't"
}


def simple_tokenize(text: str) -> List[str]:
    """简单的文本分词函数"""
    # 转小写并用正则表达式分词
    text = text.lower()
    # 提取所有字母组成的单词
    tokens = re.findall(r'\b[a-z]+\b', text)
    return tokens


def extract_ngrams(tokens: List[str], n: int) -> List[str]:
    """提取n-gram"""
    if len(tokens) < n:
        return []
    ngrams_list = []
    for i in range(len(tokens) - n + 1):
        ngram = '_'.join(tokens[i:i+n])
        ngrams_list.append(ngram)
    return ngrams_list


class NGramFeatureExtractor:
    """N-gram特征提取器"""
    
    def __init__(self, n_range=(1, 3), max_features=10000, min_freq=2):
        """
        初始化特征提取器
        
        Args:
            n_range: n-gram的范围，例如(1,3)表示1-gram到3-gram
            max_features: 最大特征数量
            min_freq: 最小出现频率
        """
        self.n_range = n_range
        self.max_features = max_features
        self.min_freq = min_freq
        self.feature_to_idx = {}
        self.idx_to_feature = {}
        self.vocab_size = 0
        self.stop_words = ENGLISH_STOP_WORDS
    
    def preprocess_text(self, text: str) -> List[str]:
        """文本预处理"""
        # 使用简单分词
        tokens = simple_tokenize(text)
        # 去停用词和短词
        tokens = [w for w in tokens if w not in self.stop_words and len(w) > 1]
        return tokens
    
    def extract_ngrams(self, tokens: List[str]) -> List[str]:
        """提取n-gram特征"""
        all_ngrams = []
        for n in range(self.n_range[0], self.n_range[1] + 1):
            n_grams = extract_ngrams(tokens, n)
            all_ngrams.extend(n_grams)
        return all_ngrams
    
    def build_vocabulary(self, texts: List[str]):
        """构建词汇表"""
        print("正在构建词汇表...")
        ngram_counts = Counter()
        
        for text in texts:
            tokens = self.preprocess_text(text)
            text_ngrams = self.extract_ngrams(tokens)
            ngram_counts.update(text_ngrams)
        
        # 过滤低频词并选择top特征
        filtered_ngrams = {ngram: count for ngram, count in ngram_counts.items() 
                          if count >= self.min_freq}
        
        # 按频率排序，取前max_features个
        top_ngrams = sorted(filtered_ngrams.items(), key=lambda x: x[1], reverse=True)
        top_ngrams = top_ngrams[:self.max_features]
        
        # 构建特征映射
        self.feature_to_idx = {ngram: idx for idx, (ngram, _) in enumerate(top_ngrams)}
        self.idx_to_feature = {idx: ngram for ngram, idx in self.feature_to_idx.items()}
        self.vocab_size = len(self.feature_to_idx)
        
        print(f"词汇表大小: {self.vocab_size}")
        print(f"前10个特征: {list(self.feature_to_idx.keys())[:10]}")
    
    def text_to_features(self, text: str) -> np.ndarray:
        """将文本转换为特征向量"""
        tokens = self.preprocess_text(text)
        text_ngrams = self.extract_ngrams(tokens)
        
        # 创建特征向量
        feature_vector = np.zeros(self.vocab_size)
        ngram_counts = Counter(text_ngrams)
        
        for ngram, count in ngram_counts.items():
            if ngram in self.feature_to_idx:
                idx = self.feature_to_idx[ngram]
                feature_vector[idx] = count
        
        return feature_vector
    
    def texts_to_features(self, texts: List[str]) -> np.ndarray:
        """批量转换文本为特征矩阵"""
        print("正在提取特征...")
        features = []
        for i, text in enumerate(texts):
            if i % 1000 == 0:
                print(f"处理进度: {i}/{len(texts)}")
            features.append(self.text_to_features(text))
        return np.array(features)


class MaxEntClassifier:
    """最大熵分类器"""
    
    def __init__(self, num_classes: int, feature_dim: int, learning_rate=0.01, 
                 regularization=0.01):
        """
        初始化最大熵分类器
        
        Args:
            num_classes: 类别数量
            feature_dim: 特征维度
            learning_rate: 学习率
            regularization: 正则化系数
        """
        self.num_classes = num_classes
        self.feature_dim = feature_dim
        self.learning_rate = learning_rate
        self.regularization = regularization
        
        # 初始化权重矩阵 (feature_dim x num_classes)
        self.weights = np.random.normal(0, 0.01, (feature_dim, num_classes))
        self.bias = np.zeros(num_classes)
        
        # 训练历史
        self.loss_history = []
        self.accuracy_history = []
    
    def compute_scores(self, X: np.ndarray) -> np.ndarray:
        """计算分类得分"""
        # scores = X @ weights + bias
        scores = np.dot(X, self.weights) + self.bias
        return scores
    
    def softmax(self, scores: np.ndarray) -> np.ndarray:
        """计算softmax概率"""
        # 数值稳定性处理
        scores = scores - np.max(scores, axis=1, keepdims=True)
        exp_scores = np.exp(scores)
        probabilities = exp_scores / np.sum(exp_scores, axis=1, keepdims=True)
        return probabilities
    
    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        """预测概率"""
        scores = self.compute_scores(X)
        return self.softmax(scores)
    
    def predict(self, X: np.ndarray) -> np.ndarray:
        """预测类别"""
        probabilities = self.predict_proba(X)
        return np.argmax(probabilities, axis=1)
    
    def compute_loss(self, X: np.ndarray, y: np.ndarray) -> float:
        """计算最大熵损失（交叉熵损失）"""
        n_samples = X.shape[0]
        
        # 计算概率
        probabilities = self.predict_proba(X)
        
        # 计算交叉熵损失
        # 避免log(0)
        probabilities = np.clip(probabilities, 1e-15, 1 - 1e-15)
        cross_entropy = -np.mean(np.log(probabilities[np.arange(n_samples), y]))
        
        # 添加L2正则化
        l2_penalty = self.regularization * np.sum(self.weights ** 2)
        
        total_loss = cross_entropy + l2_penalty
        return total_loss
    
    def compute_gradients(self, X: np.ndarray, y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """计算梯度"""
        n_samples = X.shape[0]
        
        # 前向传播
        probabilities = self.predict_proba(X)
        
        # 创建one-hot编码的真实标签
        y_one_hot = np.zeros((n_samples, self.num_classes))
        y_one_hot[np.arange(n_samples), y] = 1
        
        # 计算梯度
        # dL/dW = X^T @ (P - Y) / n + 2 * lambda * W
        prob_diff = probabilities - y_one_hot
        weights_grad = np.dot(X.T, prob_diff) / n_samples + 2 * self.regularization * self.weights
        bias_grad = np.mean(prob_diff, axis=0)
        
        return weights_grad, bias_grad
    
    def fit(self, X: np.ndarray, y: np.ndarray, X_val=None, y_val=None, 
            epochs=100, batch_size=32, verbose=True):
        """训练模型"""
        n_samples = X.shape[0]
        n_batches = (n_samples + batch_size - 1) // batch_size
        
        print(f"开始训练，样本数: {n_samples}, 批次数: {n_batches}, 轮数: {epochs}")
        
        for epoch in range(epochs):
            # 随机打乱数据
            indices = np.random.permutation(n_samples)
            X_shuffled = X[indices]
            y_shuffled = y[indices]
            
            epoch_loss = 0
            
            # 小批次梯度下降
            for batch_idx in range(n_batches):
                start_idx = batch_idx * batch_size
                end_idx = min((batch_idx + 1) * batch_size, n_samples)
                
                X_batch = X_shuffled[start_idx:end_idx]
                y_batch = y_shuffled[start_idx:end_idx]
                
                # 计算梯度
                weights_grad, bias_grad = self.compute_gradients(X_batch, y_batch)
                
                # 更新参数
                self.weights -= self.learning_rate * weights_grad
                self.bias -= self.learning_rate * bias_grad
                
                # 累计损失
                batch_loss = self.compute_loss(X_batch, y_batch)
                epoch_loss += batch_loss * len(X_batch)
            
            # 计算平均损失
            avg_loss = epoch_loss / n_samples
            self.loss_history.append(avg_loss)
            
            # 计算准确率
            train_accuracy = self.accuracy(X, y)
            self.accuracy_history.append(train_accuracy)
            
            if verbose and (epoch + 1) % 10 == 0:
                print(f"Epoch {epoch + 1}/{epochs}, Loss: {avg_loss:.4f}, "
                      f"Train Accuracy: {train_accuracy:.4f}")
                
                if X_val is not None and y_val is not None:
                    val_accuracy = self.accuracy(X_val, y_val)
                    print(f"                           Val Accuracy: {val_accuracy:.4f}")
    
    def accuracy(self, X: np.ndarray, y: np.ndarray) -> float:
        """计算准确率"""
        predictions = self.predict(X)
        return np.mean(predictions == y)


class ModelEvaluator:
    """模型评估器"""
    
    @staticmethod
    def plot_training_history(classifier: MaxEntClassifier):
        """绘制训练历史"""
        try:
            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
            
            # 损失曲线
            ax1.plot(classifier.loss_history)
            ax1.set_title('Training Loss')
            ax1.set_xlabel('Epoch')
            ax1.set_ylabel('Loss')
            ax1.grid(True)
            
            # 准确率曲线
            ax2.plot(classifier.accuracy_history)
            ax2.set_title('Training Accuracy')
            ax2.set_xlabel('Epoch')
            ax2.set_ylabel('Accuracy')
            ax2.grid(True)
            
            plt.tight_layout()
            plt.show()
        except Exception as e:
            print(f"无法显示图表: {e}")
    
    @staticmethod
    def evaluate_model(classifier: MaxEntClassifier, X_test: np.ndarray, 
                      y_test: np.ndarray, class_names: List[str] = None):
        """全面评估模型"""
        predictions = classifier.predict(X_test)
        probabilities = classifier.predict_proba(X_test)
        
        accuracy = np.mean(predictions == y_test)
        print(f"测试准确率: {accuracy:.4f}")
        
        # 计算每个类别的精确率和召回率
        num_classes = len(np.unique(y_test))
        if class_names is None:
            class_names = [f"Class {i}" for i in range(num_classes)]
        
        print("\n各类别详细指标:")
        print("-" * 50)
        
        for i in range(num_classes):
            true_positive = np.sum((predictions == i) & (y_test == i))
            false_positive = np.sum((predictions == i) & (y_test != i))
            false_negative = np.sum((predictions != i) & (y_test == i))
            
            precision = true_positive / (true_positive + false_positive) if (true_positive + false_positive) > 0 else 0
            recall = true_positive / (true_positive + false_negative) if (true_positive + false_negative) > 0 else 0
            f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
            
            print(f"{class_names[i]}: Precision={precision:.4f}, Recall={recall:.4f}, F1={f1:.4f}")
        
        return accuracy


def load_sst5_dataset(data_dir: str):
    """加载SST-5数据集"""
    import json
    import os
    
    print("加载SST-5数据集...")
    
    def load_jsonl(file_path):
        texts = []
        labels = []
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                data = json.loads(line.strip())
                texts.append(data['text'])
                labels.append(data['label'])
        return texts, labels
    
    # 加载训练、验证和测试集
    train_texts, train_labels = load_jsonl(os.path.join(data_dir, 'train.jsonl'))
    dev_texts, dev_labels = load_jsonl(os.path.join(data_dir, 'dev.jsonl'))
    test_texts, test_labels = load_jsonl(os.path.join(data_dir, 'test.jsonl'))
    
    print(f"训练集: {len(train_texts)} 样本")
    print(f"验证集: {len(dev_texts)} 样本")
    print(f"测试集: {len(test_texts)} 样本")
    
    # 打印标签分布
    from collections import Counter
    print(f"训练集标签分布: {Counter(train_labels)}")
    print(f"验证集标签分布: {Counter(dev_labels)}")
    print(f"测试集标签分布: {Counter(test_labels)}")
    
    return (train_texts, train_labels), (dev_texts, dev_labels), (test_texts, test_labels)


def create_demo_dataset():
    """创建演示数据集（保留用于测试）"""
    print("创建演示数据集...")
    
    # 简单的情感分析数据集
    positive_texts = [
        "I love this movie it's amazing and wonderful",
        "This is the best film I have ever seen fantastic",
        "Excellent acting and great story highly recommended",
        "Beautiful cinematography and outstanding performance",
        "Absolutely brilliant work of art love it",
        "Perfect movie with great direction and script",
        "Amazing storyline and excellent character development",
        "Outstanding film with beautiful visuals and music",
        "Incredible performance by all actors wonderful",
        "Magnificent movie experience truly enjoyed it",
        "Superb acting and direction really impressed me",
        "Fantastic movie with great plot and characters",
        "Wonderful film highly recommend to everyone",
        "Excellent cinematography and brilliant acting",
        "Amazing story with outstanding visual effects"
    ] * 15  # 扩展数据集
    
    negative_texts = [
        "This movie is terrible and boring waste of time",
        "Awful acting and poor story completely disappointed",
        "Worst film ever made absolutely horrible",
        "Bad direction and terrible script not recommended",
        "Completely boring and uninteresting movie",
        "Poor performance by actors and weak plot",
        "Disappointing movie with bad cinematography",
        "Terrible waste of time and money avoid it",
        "Awful storyline and poor character development",
        "Horrible movie experience completely disappointed",
        "Terrible acting and boring plot disaster",
        "Worst movie ever seen complete waste time",
        "Awful film with terrible direction and script",
        "Poor quality movie with bad acting",
        "Disappointing film avoid at all costs"
    ] * 15  # 扩展数据集
    
    neutral_texts = [
        "The movie was okay nothing special but watchable",
        "Average film with decent acting and story",
        "It's an ordinary movie not bad not great",
        "The film has some good moments but overall average",
        "Acceptable movie with reasonable story and acting",
        "Not the best but not the worst either",
        "Average performance by actors and decent plot",
        "Ordinary film with standard cinematography",
        "The movie is fine but nothing outstanding",
        "Decent movie experience neither good nor bad",
        "Acceptable film with average story and acting",
        "Standard movie nothing particularly special",
        "Ordinary film with decent but not great plot",
        "Average movie experience could be better",
        "Fair movie with reasonable story development"
    ] * 15  # 扩展数据集
    
    # 合并数据
    texts = positive_texts + negative_texts + neutral_texts
    labels = [1] * len(positive_texts) + [0] * len(negative_texts) + [2] * len(neutral_texts)
    
    # 随机打乱
    combined = list(zip(texts, labels))
    random.shuffle(combined)
    texts, labels = zip(*combined)
    
    return list(texts), list(labels)


def hyperparameter_tuning(X_train, y_train, X_val, y_val, feature_extractor):
    """超参数调优"""
    print("开始超参数调优...")
    
    # 定义超参数网格
    learning_rates = [0.001, 0.01, 0.1]
    regularizations = [0.001, 0.01, 0.1]
    
    best_accuracy = 0
    best_params = {}
    results = []
    
    for lr in learning_rates:
        for reg in regularizations:
            print(f"测试 lr={lr}, reg={reg}")
            
            # 创建和训练模型
            classifier = MaxEntClassifier(
                num_classes=5,  # 修改为5分类
                feature_dim=feature_extractor.vocab_size,
                learning_rate=lr,
                regularization=reg
            )
            
            classifier.fit(X_train, y_train, X_val, y_val, epochs=50, verbose=False)
            
            # 评估验证集
            val_accuracy = classifier.accuracy(X_val, y_val)
            results.append((lr, reg, val_accuracy))
            
            if val_accuracy > best_accuracy:
                best_accuracy = val_accuracy
                best_params = {'learning_rate': lr, 'regularization': reg}
            
            print(f"验证准确率: {val_accuracy:.4f}")
    
    print(f"\n最佳参数: {best_params}")
    print(f"最佳验证准确率: {best_accuracy:.4f}")
    
    return best_params, results


def main():
    """主函数 - 完整的机器学习流程演示"""
    print("=" * 60)
    print("基于N-gram的最大熵分类器 - SST-5五分类任务")
    print("=" * 60)
    
    # 1. 加载SST-5数据集
    sst5_data_dir = "nlp/dataset/sst5"
    (train_texts, train_labels), (dev_texts, dev_labels), (test_texts, test_labels) = load_sst5_dataset(sst5_data_dir)
    
    class_names = ['Very Negative', 'Negative', 'Neutral', 'Positive', 'Very Positive']
    
    print(f"数据集大小: 训练集{len(train_texts)}, 验证集{len(dev_texts)}, 测试集{len(test_texts)}")
    
    # 2. 特征提取
    feature_extractor = NGramFeatureExtractor(
        n_range=(1, 3),
        max_features=10000,  # 增加特征数量以处理更大的数据集
        min_freq=3  # 提高最小频率以减少噪声
    )
    
    # 构建词汇表（只用训练集）
    feature_extractor.build_vocabulary(train_texts)
    
    # 转换为特征矩阵
    X_train = feature_extractor.texts_to_features(train_texts)
    X_val = feature_extractor.texts_to_features(dev_texts)
    X_test = feature_extractor.texts_to_features(test_texts)
    
    y_train = np.array(train_labels)
    y_val = np.array(dev_labels)
    y_test = np.array(test_labels)
    
    print(f"特征矩阵形状: {X_train.shape}")
    
    # 3. 超参数调优
    best_params, tuning_results = hyperparameter_tuning(X_train, y_train, X_val, y_val, feature_extractor)
    
    # 4. 使用最佳参数训练最终模型
    print("\n使用最佳参数训练最终模型...")
    final_classifier = MaxEntClassifier(
        num_classes=5,  # 修改为5分类
        feature_dim=feature_extractor.vocab_size,
        learning_rate=best_params['learning_rate'],
        regularization=best_params['regularization']
    )
    
    final_classifier.fit(X_train, y_train, X_val, y_val, epochs=100, verbose=True)
    
    # 5. 模型评估
    print("\n" + "=" * 50)
    print("最终模型评估")
    print("=" * 50)
    
    evaluator = ModelEvaluator()
    test_accuracy = evaluator.evaluate_model(final_classifier, X_test, y_test, class_names)
    
    # 6. 可视化训练过程
    evaluator.plot_training_history(final_classifier)
    
    # 7. 演示预测
    print("\n" + "=" * 50)
    print("预测演示")
    print("=" * 50)
    
    demo_texts = [
        "This movie is absolutely fantastic and amazing",
        "The film was terrible and completely boring",
        "It's an okay movie nothing too special",
        "One of the worst films I have ever seen",
        "A masterpiece of cinema truly outstanding"
    ]
    
    for text in demo_texts:
        features = feature_extractor.text_to_features(text)
        prediction = final_classifier.predict(features.reshape(1, -1))[0]
        probabilities = final_classifier.predict_proba(features.reshape(1, -1))[0]
        
        print(f"文本: {text}")
        print(f"预测类别: {class_names[prediction]}")
        print(f"概率分布: {dict(zip(class_names, probabilities))}")
        print("-" * 50)
    
    print("\n训练完成！现在你可以体验以下核心功能:")
    print("1. N-gram特征提取: feature_extractor.extract_ngrams()")
    print("2. 特征向量转换: feature_extractor.text_to_features()")
    print("3. 分类得分计算: final_classifier.compute_scores()")
    print("4. Softmax概率计算: final_classifier.softmax()")
    print("5. 损失函数计算: final_classifier.compute_loss()")
    print("6. 梯度计算: final_classifier.compute_gradients()")
    print("7. 模型预测: final_classifier.predict()")
    
    return final_classifier, feature_extractor


if __name__ == "__main__":
    model, extractor = main()